from langdev_helper.util.broswer import Broswer
from langdev_helper.util.search import get_link_content, print_bing_search_result, print_bing_result_lite
from langgraph.prebuilt import ToolNode, tools_condition

def continue_msg(cmd_content):
    if cmd_content.startswith('/cwd'):
        cmd_content = '当前目录'
    elif cmd_content.startswith('/list'):
        cmd_content = 'show cwd and list dir: cwd, output format use markdown'
    else:
        return None
    return {"messages": [{"role": "user", "content": cmd_content}]}

def end_msg(cmd_content):
    if cmd_content.startswith('/h'):
        print('''
help:

- direct
    /h     help
    /bing  bing search
    /blist bing list

- llm_tool
    /cwd   current working directory
    /list  list dir
        ''')
        pass
    elif cmd_content.startswith('/bing'):
        print('\nkeyword:', Broswer.keyword, cmd_content, '\n')
        
    elif cmd_content.startswith('/bl') or cmd_content.startswith('/blist'):
        print('\nresults :')
        print_bing_result_lite(Broswer.results)
    elif cmd_content.startswith('/bv'):
        params = cmd_content.split(' ')
        link_num = params[1]
        
        item = Broswer.results[int(link_num)-1]
        print('\ntitle:', item['title'], '  ',item['link']
        # , '\nlink :', item['link']
        )

        content = get_link_content(item['link'])
        clen = len(content)
        print(f'\ncontent[{clen}]:', content[0: 850])
    else:
        return None

    return {"messages": [{"role": "ai", "content": 'no'}]}


class CmdNode:
    """A node that runs the tools requested in the last AIMessage."""

    def __init__(self) -> None:
        pass

    def __call__(self, inputs: dict):
        if messages := inputs.get("messages", []):
            message = messages[-1]
        else:
            raise ValueError("No message found in input")
        
        # print('  <== CmdNode inputs:', inputs)

        last_message = messages[-1]
        # print('last_message', last_message)

        usr_message = messages[0]
        cmd_content = last_message.content

        if cmd_content.startswith('/'):
            # print('\ncmd:', cmd_content)
            ret = continue_msg(cmd_content)
            if not ret:
                ret = end_msg(cmd_content)
            return ret
        else:
            return {"messages": [{"role": "user", "content": cmd_content}]} 


from typing import Annotated
from langgraph.graph.message import add_messages
from langgraph.graph import StateGraph, START, END
from typing_extensions import TypedDict

class State(TypedDict):
    messages: Annotated[list, add_messages]

def cmd_condition(
    state: State,
):

    if isinstance(state, list):
        ai_message = state[-1]
    elif messages := state.get("messages", []):
        ai_message = messages[-1]
    else:
        raise ValueError(f"No messages found in input state to tool_edge: {state}")
    
    # print('cmd_condition:', ai_message)
    # print('=== ai_message type: ', isinstance(ai_message, HumanMessage))
    # ai_message.pretty_print()
    if(ai_message.content == 'no'):
        return END
    
    return "AGENT"

def sense_tools_condition(
    state: State,
):

    if isinstance(state, list):
        ai_message = state[-1]
    elif messages := state.get("messages", []):
        ai_message = messages[-1]
    else:
        raise ValueError(f"No messages found in input state to tool_edge: {state}")
    
    # print('cmd_condition:', ai_message)
    # print('=== ai_message type: ', isinstance(ai_message, HumanMessage))
    # ai_message.pretty_print()
    if(ai_message.content == 'no'):
        return END
    
    return "AGENT"

def route_tools(state: State):
    next_node = tools_condition(state)
    # If no tools are invoked, return to the user
    if next_node == END:
        return END
    ai_message = state["messages"][-1]
    # This assumes single tool calls. To handle parallel tool calling, you'd want to
    # use an ANY condition
    first_tool_call = ai_message.tool_calls[0]
    if first_tool_call["name"] in ['CwdTool', 'WriteFileTool']:
        return "sense_tools"
    return "tools"